#include "../include/sm4.h"
#include <time.h>

// SM4宏定义
#define SM4_ROTL32(x, n) (((x) << (n)) | ((x) >> (32 - (n))))

// SM4常量表
static const uint32_t SM4_CK[32] = {
    0x00070e15, 0x1c232a31, 0x383f464d, 0x545b6269,
    0x70777e85, 0x8c939aa1, 0xa8afb6bd, 0xc4cbd2d9,
    0xe0e7eef5, 0xfc030a11, 0x181f262d, 0x343b4249,
    0x50575e65, 0x6c737a81, 0x888f969d, 0xa4abb2b9,
    0xc0c7ced5, 0xdce3eaf1, 0xf8ff060d, 0x141b2229,
    0x30373e45, 0x4c535a61, 0x686f767d, 0x848b9299,
    0xa0a7aeb5, 0xbcc3cad1, 0xd8dfe6ed, 0xf4fb0209,
    0x10171e25, 0x2c333a41, 0x484f565d, 0x646b7279
};

static const uint32_t SM4_FK[4] = {
    0xA3B1BAC6, 0x56AA3350, 0x677D9197, 0xB27022DC
};

static const uint8_t SM4_SBOX[256] = {
    0xd6,0x90,0xe9,0xfe,0xcc,0xe1,0x3d,0xb7,0x16,0xb6,0x14,0xc2,0x28,0xfb,0x2c,0x05,
    0x2b,0x67,0x9a,0x76,0x2a,0xbe,0x04,0xc3,0xaa,0x44,0x13,0x26,0x49,0x86,0x06,0x99,
    0x9c,0x42,0x50,0xf4,0x91,0xef,0x98,0x7a,0x33,0x54,0x0b,0x43,0xed,0xcf,0xac,0x62,
    0xe4,0xb3,0x1c,0xa9,0xc9,0x08,0xe8,0x95,0x80,0xdf,0x94,0xfa,0x75,0x8f,0x3f,0xa6,
    0x47,0x07,0xa7,0xfc,0xf3,0x73,0x17,0xba,0x83,0x59,0x3c,0x19,0xe6,0x85,0x4f,0xa8,
    0x68,0x6b,0x81,0xb2,0x71,0x64,0xda,0x8b,0xf8,0xeb,0x0f,0x4b,0x70,0x56,0x9d,0x35,
    0x1e,0x24,0x0e,0x5e,0x63,0x58,0xd1,0xa2,0x25,0x22,0x7c,0x3b,0x01,0x21,0x78,0x87,
    0xd4,0x00,0x46,0x57,0x9f,0xd3,0x27,0x52,0x4c,0x36,0x02,0xe7,0xa0,0xc4,0xc8,0x9e,
    0xea,0xbf,0x8a,0xd2,0x40,0xc7,0x38,0xb5,0xa3,0xf7,0xf2,0xce,0xf9,0x61,0x15,0xa1,
    0xe0,0xae,0x5d,0xa4,0x9b,0x34,0x1a,0x55,0xad,0x93,0x32,0x30,0xf5,0x8c,0xb1,0xe3,
    0x1d,0xf6,0xe2,0x2e,0x82,0x66,0xca,0x60,0xc0,0x29,0x23,0xab,0x0d,0x53,0x4e,0x6f,
    0xd5,0xdb,0x37,0x45,0xde,0xfd,0x8e,0x2f,0x03,0xff,0x6a,0x72,0x6d,0x6c,0x5b,0x51,
    0x8d,0x1b,0xaf,0x92,0xbb,0xdd,0xbc,0x7f,0x11,0xd9,0x5c,0x41,0x1f,0x10,0x5a,0xd8,
    0x0a,0xc1,0x31,0x88,0xa5,0xcd,0x7b,0xbd,0x2d,0x74,0xd0,0x12,0xb8,0xe5,0xb4,0xb0,
    0x89,0x69,0x97,0x4a,0x0c,0x96,0x77,0x7e,0x65,0xb9,0xf1,0x09,0xc5,0x6e,0xc6,0x84,
    0x18,0xf0,0x7d,0xec,0x3a,0xdc,0x4d,0x20,0x79,0xee,0x5f,0x3e,0xd7,0xcb,0x39,0x48
};

// 工具函数
void sm4_print_hex(const char *label, const uint8_t *data, size_t len) {
    printf("%s: ", label);
    for (size_t i = 0; i < len; i++) {
        printf("%02X", data[i]);
    }
    printf("\n");
}

void sm4_generate_key(uint8_t key[SM4_KEY_SIZE]) {
    srand(time(NULL));
    for (int i = 0; i < SM4_KEY_SIZE; i++) {
        key[i] = rand() & 0xFF;
    }
}

// 密钥扩展
static void sm4_key_schedule(const uint8_t key[SM4_KEY_SIZE], uint32_t rk[SM4_ROUNDS]) {
    uint32_t K[36];
    int i;
    
    // 初始密钥加
    for (i = 0; i < 4; i++) {
        K[i] = SM4_FK[i] ^ ((key[4 * i] << 24) | (key[4 * i + 1] << 16) |
                         (key[4 * i + 2] << 8) | (key[4 * i + 3]));
    }
    
    // 轮密钥生成
    for (i = 0; i < SM4_ROUNDS; i++) {
        uint32_t tmp = K[i + 1] ^ K[i + 2] ^ K[i + 3] ^ SM4_CK[i];
        
        // S盒变换
        uint32_t sbox_out = (SM4_SBOX[(tmp >> 24) & 0xFF] << 24) |
                          (SM4_SBOX[(tmp >> 16) & 0xFF] << 16) |
                          (SM4_SBOX[(tmp >> 8) & 0xFF] << 8) |
                          (SM4_SBOX[tmp & 0xFF]);
        
        // 线性变换
        K[i + 4] = K[i] ^ (sbox_out ^ SM4_ROTL32(sbox_out, 13) ^ SM4_ROTL32(sbox_out, 23));
        rk[i] = K[i + 4];
    }
}

// 轮函数
static uint32_t sm4_round_function(uint32_t x0, uint32_t x1, uint32_t x2, uint32_t x3, uint32_t rk) {
    uint32_t tmp = x1 ^ x2 ^ x3 ^ rk;
    
    // S盒变换
    uint32_t sbox_out = (SM4_SBOX[(tmp >> 24) & 0xFF] << 24) |
                      (SM4_SBOX[(tmp >> 16) & 0xFF] << 16) |
                      (SM4_SBOX[(tmp >> 8) & 0xFF] << 8) |
                      (SM4_SBOX[tmp & 0xFF]);
    
    // 线性变换
    return x0 ^ sbox_out ^ SM4_ROTL32(sbox_out, 2) ^ SM4_ROTL32(sbox_out, 10) ^
           SM4_ROTL32(sbox_out, 18) ^ SM4_ROTL32(sbox_out, 24);
}

// SM4加密
void sm4_encrypt(const uint8_t key[SM4_KEY_SIZE], 
                 const uint8_t plaintext[SM4_BLOCK_SIZE], 
                 uint8_t ciphertext[SM4_BLOCK_SIZE]) {
    uint32_t rk[SM4_ROUNDS];
    uint32_t X[36];
    int i;
    
    // 密钥扩展
    sm4_key_schedule(key, rk);
    
    // 初始轮输入
    for (i = 0; i < 4; i++) {
        X[i] = (plaintext[4 * i] << 24) | (plaintext[4 * i + 1] << 16) |
               (plaintext[4 * i + 2] << 8) | (plaintext[4 * i + 3]);
    }
    
    // 32轮迭代
    for (i = 0; i < SM4_ROUNDS; i++) {
        X[i + 4] = sm4_round_function(X[i], X[i + 1], X[i + 2], X[i + 3], rk[i]);
    }
    
    // 反序变换输出
    for (i = 0; i < 4; i++) {
        ciphertext[4 * i] = (X[35 - i] >> 24) & 0xFF;
        ciphertext[4 * i + 1] = (X[35 - i] >> 16) & 0xFF;
        ciphertext[4 * i + 2] = (X[35 - i] >> 8) & 0xFF;
        ciphertext[4 * i + 3] = X[35 - i] & 0xFF;
    }
}

// SM4解密
void sm4_decrypt(const uint8_t key[SM4_KEY_SIZE], 
                 const uint8_t ciphertext[SM4_BLOCK_SIZE], 
                 uint8_t plaintext[SM4_BLOCK_SIZE]) {
    uint32_t rk[SM4_ROUNDS];
    uint32_t X[36];
    int i;
    
    // 密钥扩展
    sm4_key_schedule(key, rk);
    
    // 初始轮输入
    for (i = 0; i < 4; i++) {
        X[i] = (ciphertext[4 * i] << 24) | (ciphertext[4 * i + 1] << 16) |
               (ciphertext[4 * i + 2] << 8) | (ciphertext[4 * i + 3]);
    }
    
    // 32轮迭代（使用逆序轮密钥）
    for (i = 0; i < SM4_ROUNDS; i++) {
        X[i + 4] = sm4_round_function(X[i], X[i + 1], X[i + 2], X[i + 3], rk[SM4_ROUNDS - 1 - i]);
    }
    
    // 反序变换输出
    for (i = 0; i < 4; i++) {
        plaintext[4 * i] = (X[35 - i] >> 24) & 0xFF;
        plaintext[4 * i + 1] = (X[35 - i] >> 16) & 0xFF;
        plaintext[4 * i + 2] = (X[35 - i] >> 8) & 0xFF;
        plaintext[4 * i + 3] = X[35 - i] & 0xFF;
    }
}

// ECB模式加密（带PKCS7填充）
void sm4_ecb_encrypt(const uint8_t key[SM4_KEY_SIZE], 
                     const uint8_t *input, 
                     uint8_t *output, 
                     size_t length) {
    size_t blocks = length / SM4_BLOCK_SIZE;
    size_t remaining = length % SM4_BLOCK_SIZE;
    
    // 加密完整的数据块
    for (size_t i = 0; i < blocks; i++) {
        sm4_encrypt(key, input + i * SM4_BLOCK_SIZE, output + i * SM4_BLOCK_SIZE);
    }
    
    // 处理最后一个不完整的块
    if (remaining > 0) {
        uint8_t last_block[SM4_BLOCK_SIZE] = {0};
        // 复制剩余数据
        memcpy(last_block, input + blocks * SM4_BLOCK_SIZE, remaining);
        // PKCS7填充
        uint8_t pad_value = SM4_BLOCK_SIZE - remaining;
        for (size_t i = remaining; i < SM4_BLOCK_SIZE; i++) {
            last_block[i] = pad_value;
        }
        // 加密填充后的块
        sm4_encrypt(key, last_block, output + blocks * SM4_BLOCK_SIZE);
    }
}

// ECB模式解密（带PKCS7去除填充）
void sm4_ecb_decrypt(const uint8_t key[SM4_KEY_SIZE], 
                     const uint8_t *input, 
                     uint8_t *output, 
                     size_t length) {
    // 计算实际的数据块数（包括填充块）
    size_t total_blocks = (length + SM4_BLOCK_SIZE - 1) / SM4_BLOCK_SIZE;
    
    // 解密所有数据块
    for (size_t i = 0; i < total_blocks; i++) {
        sm4_decrypt(key, input + i * SM4_BLOCK_SIZE, output + i * SM4_BLOCK_SIZE);
    }
}
// 自检测试
int sm4_self_test(void) {
    printf("=== SM4 自检测试 ===\n");
    
    // 标准测试向量
    uint8_t key[16] = {
        0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF,
        0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32, 0x10
    };
    
    uint8_t plaintext[16] = {
        0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF,
        0xFE, 0xDC, 0xBA, 0x98, 0x76, 0x54, 0x32, 0x10
    };
    
    uint8_t expected_cipher[16] = {
        0x68, 0x1E, 0xDF, 0x34, 0xD2, 0x06, 0x96, 0x5E,
        0x86, 0xB3, 0xE9, 0x4F, 0x53, 0x6E, 0x42, 0x46
    };
    
    uint8_t ciphertext[16];
    uint8_t decrypted[16];
    
    // 加密测试
    sm4_encrypt(key, plaintext, ciphertext);
    
    sm4_print_hex("密钥", key, 16);
    sm4_print_hex("明文", plaintext, 16);
    sm4_print_hex("加密结果", ciphertext, 16);
    sm4_print_hex("期望密文", expected_cipher, 16);
    
    // 解密测试
    sm4_decrypt(key, expected_cipher, decrypted);
    sm4_print_hex("解密结果", decrypted, 16);
    
    // 验证结果
    if (memcmp(ciphertext, expected_cipher, 16) == 0 && 
        memcmp(decrypted, plaintext, 16) == 0) {
        printf("✓ SM4 自检测试通过！\n\n");
        return 0;
    } else {
        printf("✗ SM4 自检测试失败！\n\n");
        return 1;
    }
}
